package cn.bitkit.tools.crypto;

import cn.bitkit.base.util.StringUtil;
import lombok.Cleanup;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.security.*;

/**
 * 
 * @author changbo
 *
 */
@Slf4j
public class RSA {
	
    public static final String ALGORITHM = "RSA";
	public static final String SIGNATURE_ALGORITHM = "SHA256withRSA";
    public static final String CIPHER_ALGORITHM = "RSA/ECB/PKCS1Padding";
    
    /** 分段加密长度,RSA最大加密明文大小max = KEY_SIZE/8 - 11 */
    public static final int MAX_ENCRYPT_BLOCK = 245;;
    /** 分段解密长度,RSA最大解密密文大小max = KEY_SIZE/8 */
    public static final int MAX_DECRYPT_BLOCK = 256;

    public static KeyMap createKeyMap(){
        return createKeyMap(2048);
    }

    @SneakyThrows
    public static KeyMap createKeyMap(int keySize){
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM);
        keyPairGenerator.initialize(keySize, new SecureRandom());
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        return new KeyMap(keyPair);
    }
    
    public static String encrypt(String data, String publicKey){
    	return encrypt(data, publicKey, CIPHER_ALGORITHM);
    }
    
    public static String encrypt(String data, String publicKey, String cipherAlgorithm){
        PublicKey pubKey = CertUtil.getEncodedPublicKey(publicKey);
        byte[] res = cipher(data.getBytes(), pubKey, cipherAlgorithm, Cipher.ENCRYPT_MODE);
        return StringUtil.encodeBase64(res);
    }

    public static String decrypt(String data, String privateKey){
    	return decrypt(data, privateKey, CIPHER_ALGORITHM);
    }

    public static String decrypt(String data, String privateKey, String cipherAlgorithm){
        PrivateKey priKey = CertUtil.getEncodedPrivateKey(privateKey);
        byte[] res = cipher(StringUtil.decodeBase64(data), priKey, cipherAlgorithm, Cipher.DECRYPT_MODE);
        return new String(res);
    }

    @SneakyThrows
    public static byte[] cipher(byte[] dataBytes, Key key, String cipherAlgorithm, int mode){
        @Cleanup
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int blockSize = (mode == Cipher.ENCRYPT_MODE) ? MAX_ENCRYPT_BLOCK : MAX_DECRYPT_BLOCK;
        Cipher cipher = Cipher.getInstance(cipherAlgorithm);
        cipher.init(mode, key);
        int offSet = 0;
        int i = 0;
        // 对数据分段加/解密
        while (dataBytes.length - offSet > 0) {
            byte[] cache;
            if (dataBytes.length - offSet > blockSize) {
                cache = cipher.doFinal(dataBytes, offSet, blockSize);
            } else {
                cache = cipher.doFinal(dataBytes, offSet, dataBytes.length - offSet);
            }
            out.write(cache, 0, cache.length);
            offSet = (++i) * blockSize;
        }
        return out.toByteArray();
    }

    @SneakyThrows
    private static byte[] cipherBlock(Cipher cipher, byte[] dataBytes, int blockSize) {
        @Cleanup
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offSet = 0;
        int i = 0;
        // 对数据分段加/解密
        while (dataBytes.length - offSet > 0) {
            byte[] cache;
            if (dataBytes.length - offSet > blockSize) {
                cache = cipher.doFinal(dataBytes, offSet, blockSize);
            } else {
                cache = cipher.doFinal(dataBytes, offSet, dataBytes.length - offSet);
            }
            out.write(cache, 0, cache.length);
            offSet = (++i) * blockSize;
        }
        return out.toByteArray();
    }

	public static String sign(String data, String privateKey) {
		return sign(data, privateKey, SIGNATURE_ALGORITHM);
	}

    @SneakyThrows
	public static String sign(String data, String privateKey, String signatureAlgorithm) {
        PrivateKey priKey = CertUtil.getEncodedPrivateKey(privateKey);
        Signature signature = Signature.getInstance(signatureAlgorithm);
        signature.initSign(priKey);
        signature.update(data.getBytes());
        return StringUtil.encodeBase64(signature.sign());
	}

	public static boolean verify(String data, String publicKey, String sign) {
		return verify(data, publicKey, sign, SIGNATURE_ALGORITHM);
	}

	public static boolean verify(String data, String publicKey, String sign, String signatureAlgorithm) {
		try {
			PublicKey pubKey = CertUtil.getEncodedPublicKey(publicKey);
			Signature signature = Signature.getInstance(signatureAlgorithm);
			signature.initVerify(pubKey);
			signature.update(data.getBytes());
			return signature.verify(StringUtil.decodeBase64(sign));
	    } catch (Exception e) {
            log.error("verify:{}", e.getMessage());
	    }
        return false;
	}
}  